from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')

# 定义数据集配置文件路径
data_config = '/root/ultralytics/yolov8_face_ex01/data.yaml'  # 替换为你的data.yaml文件的实际路径

# 设置训练参数
epochs = 280
imgsz = 640
batch_size = 32
workers = 25
device = 0  # 使用GPU编号，如果使用CPU则设为None或'cpu'

# 开始训练
results = model.train(
    data=data_config,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch_size,
    workers=workers,
    device=device,
    name='yolov8_custom_train',  # 模型保存名称
    save=True,  # 是否保存模型
    exist_ok=True,  # 如果存在相同命名的训练结果是否覆盖
    cache=False  # 是否缓存数据以加速训练
)

print("训练完成:", results.save_dir)
